{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Jupyter environment detected. Enabling Open3D WebVisualizer.\n",
      "[Open3D INFO] WebRTC GUI backend enabled.\n",
      "[Open3D INFO] WebRTCWindowSystem: HTTP handshake server disabled.\n"
     ]
    }
   ],
   "source": [
    "import math\n",
    "import numpy as np\n",
    "import open3d as o3d\n",
    "import copy\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_geometries(geometries, name=\"\", show_normal=False):\n",
    "    vis = o3d.visualization.Visualizer()\n",
    "    vis.create_window(window_name=name)\n",
    "    for geometry in geometries:\n",
    "        vis.add_geometry(geometry)\n",
    "    # mesh_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(\n",
    "    #     size=20.0, origin=[0, 0, 0]\n",
    "    # )\n",
    "    # vis.add_geometry(mesh_frame)\n",
    "    opt = vis.get_render_option()\n",
    "    opt.point_show_normal = show_normal\n",
    "    opt.background_color = np.asarray([1, 1, 1])\n",
    "    opt.show_coordinate_frame = False\n",
    "    opt.line_width = 100.0\n",
    "    vis.run()\n",
    "    vis.destroy_window()\n",
    "    \n",
    "def draw_registration_result(source, target, transformation=np.identity(4), name=\"\"):\n",
    "    source_temp = copy.deepcopy(source)\n",
    "    target_temp = copy.deepcopy(target)\n",
    "    source_temp.paint_uniform_color([1, 0.706, 0])\n",
    "    target_temp.paint_uniform_color([0, 0.651, 0.929])\n",
    "    if np.linalg.norm(transformation - np.identity(4), 2) > 1e-2:\n",
    "        source_temp.transform(transformation)\n",
    "    draw_geometries([source_temp, target_temp], name)\n",
    "    \n",
    "def correspondence_connector(source, target, correspondences):\n",
    "    line_set = o3d.geometry.LineSet.create_from_point_cloud_correspondences(source, target, correspondences)\n",
    "    line_set.colors = o3d.utility.Vector3dVector([[1, 0, 0] for i in range(len(correspondences))])\n",
    "    return line_set\n",
    "\n",
    "def manual_registration(\n",
    "    source,\n",
    "    target,\n",
    "    threshold,\n",
    "    tf_init=np.identity(4),\n",
    "    manual_flag=False,\n",
    "    draw_flag=False,\n",
    "    method=\"point2point\",\n",
    "    loss=\"L2\"\n",
    "):\n",
    "    normal_flag = False\n",
    "    covariance_flag = False\n",
    "    if source.has_normals() and target.has_normals():\n",
    "        normal_flag = True\n",
    "    if source.has_covariances() and target.has_covariances():\n",
    "        covariance_flag = True\n",
    "    kernel_function = o3d.pipelines.registration.L2Loss()\n",
    "    if loss == \"Cauchy\":\n",
    "        kernel_function = o3d.pipelines.registration.CauchyLoss(2.3489)\n",
    "    elif loss == \"Huber\":\n",
    "        kernel_function = o3d.pipelines.registration.HuberLoss(1)\n",
    "    icp_method = o3d.pipelines.registration.TransformationEstimationPointToPoint(with_scaling=False)\n",
    "    if method == \"point2plane\":\n",
    "        if normal_flag:\n",
    "            icp_method = o3d.pipelines.registration.TransformationEstimationPointToPlane(kernel=kernel_function)\n",
    "    elif method == \"plane2plane\":\n",
    "        if normal_flag and covariance_flag:\n",
    "            icp_method = o3d.pipelines.registration.TransformationEstimationForGeneralizedICP(kernel=kernel_function)\n",
    "    if draw_flag:\n",
    "        print(\"Visualization of two point cloud before manual alignment\")\n",
    "        draw_registration_result(\n",
    "            source,\n",
    "            target,\n",
    "            tf_init,\n",
    "            \"Visualization of two point cloud before manual alignment\",\n",
    "        )\n",
    "    if manual_flag and np.linalg.norm(tf_init - np.identity(4), 2) < 0.01:\n",
    "        # pick points from two point clouds and builds correspondences\n",
    "        picked_id_source = pick_points(source)\n",
    "        print(\"picked source id :{}\".format(picked_id_source))\n",
    "        picked_id_target = pick_points(target)\n",
    "        print(\"picked target id :{}\".format(picked_id_target))\n",
    "        if (\n",
    "            len(picked_id_source) >= 3\n",
    "            and len(picked_id_target) >= 3\n",
    "            and len(picked_id_source) == len(picked_id_target)\n",
    "        ):\n",
    "            corr = np.zeros((len(picked_id_source), 2))\n",
    "            corr[:, 0] = picked_id_source\n",
    "            corr[:, 1] = picked_id_target\n",
    "            # estimate rough transform using the correspondences given by user\n",
    "            p2p = o3d.pipelines.registration.TransformationEstimationPointToPoint(\n",
    "                with_scaling=True\n",
    "            )\n",
    "            tf_init = p2p.compute_transformation(\n",
    "                source, target, o3d.utility.Vector2iVector(corr)\n",
    "            )\n",
    "\n",
    "    start_time = time.time()\n",
    "    icp_criteria = o3d.pipelines.registration.ICPConvergenceCriteria(relative_fitness=1e-8,relative_rmse=1e-8,max_iteration=100)\n",
    "    reg = o3d.pipelines.registration.registration_icp(\n",
    "        source,\n",
    "        target,\n",
    "        threshold,\n",
    "        init=tf_init,\n",
    "        estimation_method=icp_method,\n",
    "        criteria=icp_criteria,\n",
    "    )\n",
    "    end_time = time.time()\n",
    "    print(\"Code running time: {} s\".format(end_time - start_time))\n",
    "    print(\"Assessment criteria: {}\".format(reg))\n",
    "    # print(\"Registration result:\\n{}\".format(reg.transformation))\n",
    "    if draw_flag:\n",
    "        draw_registration_result(source, target, reg.transformation, \"Alignment result\")\n",
    "    return reg.transformation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 提取几何特征\n",
    "def preprocess_point_cloud(pcd, voxel_size):\n",
    "    # 滤除离群点\n",
    "    _, ind = pcd.remove_statistical_outlier(nb_neighbors=30, std_ratio=3.0)\n",
    "    # _, ind = pcd.remove_radius_outlier(nb_points=50,radius=2,print_progress=True)\n",
    "    pcd = pcd.select_by_index(ind, invert=False)\n",
    "    # 下采样theta = pcd.voxel_down_sample(voxel_size)\n",
    "    # 估计法线\n",
    "    pcd_down.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=voxel_size * 2, max_nn=30))\n",
    "    # 计算 FPFH 特征\n",
    "    pcd_fpfh = o3d.pipelines.registration.compute_fpfh_feature(pcd_down, o3d.geometry.KDTreeSearchParamHybrid(radius=voxel_size * 5, max_nn=100))\n",
    "    return pcd_down, pcd_fpfh\n",
    "\n",
    "# 加载点云并准备数据\n",
    "def prepare_dataset(src_points_path, tgt_points_path, voxel_size):\n",
    "    source = o3d.io.read_point_cloud(src_points_path)\n",
    "    target = o3d.io.read_point_cloud(tgt_points_path)\n",
    "    return source, target, *preprocess_point_cloud(source, voxel_size), *preprocess_point_cloud(target, voxel_size)\n",
    "\n",
    "\n",
    "def execute_corres_registration(source, target, source_feature, target_feature, voxel_size):\n",
    "    distance_threshold = voxel_size * 20\n",
    "    corres = o3d.pipelines.registration.correspondences_from_features(source_feature, target_feature, mutual_filter=True)\n",
    "    line = correspondence_connector(source, target, corres)\n",
    "    result = o3d.pipelines.registration.registration_ransac_based_on_correspondence(\n",
    "        source, target, corres,max_correspondence_distance=distance_threshold,\n",
    "        estimation_method=o3d.pipelines.registration.TransformationEstimationPointToPoint(True),\n",
    "        ransac_n=4,\n",
    "        checkers=[\n",
    "            o3d.pipelines.registration.CorrespondenceCheckerBasedOnDistance(distance_threshold),\n",
    "            o3d.pipelines.registration.CorrespondenceCheckerBasedOnNormal(0.7),\n",
    "            o3d.pipelines.registration.CorrespondenceCheckerBasedOnEdgeLength(0.9)\n",
    "        ],\n",
    "        criteria=o3d.pipelines.registration.RANSACConvergenceCriteria(confidence=1, max_iteration=4000000)\n",
    "    )\n",
    "    source_temp = copy.deepcopy(source)\n",
    "    target_temp = copy.deepcopy(target)\n",
    "    source_temp.paint_uniform_color([1, 0.706, 0])\n",
    "    target_temp.paint_uniform_color([0, 0.651, 0.929])\n",
    "    draw_geometries([source_temp.transform(result.transformation), target_temp, line])\n",
    "\n",
    "    return result\n",
    "\n",
    "def execute_global_registration(source, target, source_feature, target_feature, voxel_size):\n",
    "    result = o3d.pipelines.registration.registration_fgr_based_on_feature_matching(\n",
    "        source, target, source_feature, target_feature,\n",
    "        o3d.pipelines.registration.FastGlobalRegistrationOption(\n",
    "            maximum_correspondence_distance=20*voxel_size,\n",
    "            iteration_number=200,\n",
    "            maximum_tuple_count=1000))\n",
    "    draw_registration_result(source, target, result.transformation)\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 使用 RANSAC 进行全局配准\n",
    "pcd_data = o3d.data.DemoICPPointClouds()\n",
    "voxel_size = 0.05  # 5 cm\n",
    "# source, target, source_down, source_fpfh,target_down, target_fpfh = prepare_dataset(\"room_scan1.pcd\", \"room_scan2.pcd\", voxel_size)\n",
    "\n",
    "source, target, source_down, source_fpfh,target_down, target_fpfh = prepare_dataset(pcd_data.paths[0], pcd_data.paths[1], voxel_size)\n",
    "result = execute_global_registration(source_down,target_down,source_fpfh,target_fpfh,voxel_size)\n",
    "# manual_registration(source_down, target_down, tf_init=result.transformation, threshold=1, method=\"point2plane\", draw_flag=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Visualization of two point cloud before manual alignment\n",
      "Code running time: 0.08510303497314453 s\n",
      "Assessment criteria: RegistrationResult with fitness=1.000000e+00, inlier_rmse=1.220798e-01, and correspondence_set size of 4683\n",
      "Access transformation to get result.\n",
      "Visualization of two point cloud before manual alignment\n",
      "Code running time: 0.0919182300567627 s\n",
      "Assessment criteria: RegistrationResult with fitness=1.000000e+00, inlier_rmse=1.159516e-01, and correspondence_set size of 4683\n",
      "Access transformation to get result.\n",
      "Visualization of two point cloud before manual alignment\n",
      "Code running time: 0.08273720741271973 s\n",
      "Assessment criteria: RegistrationResult with fitness=1.000000e+00, inlier_rmse=1.155686e-01, and correspondence_set size of 4683\n",
      "Access transformation to get result.\n",
      "Visualization of two point cloud before manual alignment\n",
      "Code running time: 0.0637211799621582 s\n",
      "Assessment criteria: RegistrationResult with fitness=1.000000e+00, inlier_rmse=1.230113e-01, and correspondence_set size of 4683\n",
      "Access transformation to get result.\n"
     ]
    }
   ],
   "source": [
    "src_obb = source.get_minimal_oriented_bounding_box()\n",
    "src_obb_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(\n",
    "        size=1.0, origin=[0,0,0]\n",
    "    )\n",
    "src_obb_frame.rotate(src_obb.R, center=src_obb_frame.get_center())\n",
    "\n",
    "tgt_obb = target.get_minimal_oriented_bounding_box()\n",
    "tgt_obb_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(\n",
    "        size=1.0, origin=[0,0,0]\n",
    "    )\n",
    "tgt_obb_frame.rotate(tgt_obb.R, center=tgt_obb_frame.get_center())\n",
    "tgt_obb_frame.paint_uniform_color([0.5,0.5,0])\n",
    "\n",
    "# obb上色\n",
    "src_obb.color=[1,0,0]\n",
    "tgt_obb.color=[0,1,0]\n",
    "# 缩放\n",
    "scale = np.amax(tgt_obb.extent) / np.amax(src_obb.extent)\n",
    "# 去中心化\n",
    "source.translate(-src_obb.get_center())\n",
    "T_src_demean = np.eye(4)\n",
    "T_src_demean[:-1,-1] = -src_obb.get_center()\n",
    "target.translate(-tgt_obb.get_center())\n",
    "T_tgt_demean = np.eye(4)\n",
    "T_tgt_demean[:-1,-1] = -tgt_obb.get_center()\n",
    "src_obb.translate(-src_obb.get_center())\n",
    "tgt_obb.translate(-tgt_obb.get_center())\n",
    "\n",
    "# 估计初始旋转和缩放\n",
    "src_sorted_indices = np.argsort(src_obb.extent)\n",
    "tgt_sorted_indices = np.argsort(tgt_obb.extent)\n",
    "init_src_orient =src_obb.R[:, src_sorted_indices]\n",
    "tgt_orient =tgt_obb.R[:, tgt_sorted_indices]\n",
    "# src_obb.rotate(init_rotation, center=src_obb.get_center())\n",
    "source_down.transform(T_src_demean)\n",
    "target_down.transform(T_tgt_demean)\n",
    "# 对应obb的4种可能的坐标轴朝向\n",
    "possible_cases=np.asarray([[0., 0., 0.],\n",
    "                           [1., 0., 0.],\n",
    "                           [0., 1., 0.],\n",
    "                           [0., 0., 1.]])\n",
    "T_init = np.eye(4)\n",
    "\n",
    "for i in range(possible_cases.shape[0]):\n",
    "    src_orient = np.matmul(o3d.geometry.get_rotation_matrix_from_axis_angle(math.pi*np.matmul(init_src_orient, possible_cases[i].transpose())), init_src_orient)\n",
    "    init_rotation = np.matmul(tgt_orient, src_orient.transpose())\n",
    "    T_init[:3, :3] = scale * init_rotation\n",
    "    manual_registration(source_down, target_down, tf_init=T_init, threshold=1, method=\"point2point\", draw_flag=True)\n",
    "\n",
    "# src_obb_frame.rotate(init_rotation, center=src_obb_frame.get_center())\n",
    "# src_obb_frame.translate(tgt_obb.get_center()-src_obb.get_center())\n",
    "# source.scale(scale, center=src_obb.get_center())\n",
    "# source.rotate(init_rotation, center=src_obb.get_center())\n",
    "# source.translate(init_translation)\n",
    "# draw_geometries([src_obb, tgt_obb, src_obb_frame, tgt_obb_frame])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "vis = o3d.visualization.Visualizer()\n",
    "vis.create_window()\n",
    "theta_xyz=np.asarray([[0., 0., 0.],\n",
    "                      [1., 0., 0.],\n",
    "                      [0., 1., 0.],\n",
    "                      [0., 0., 1.]])\n",
    "T_init = np.eye(4)\n",
    "for i in range(theta_xyz.shape[0]):\n",
    "    T_init[:3, :3] = o3d.geometry.get_rotation_matrix_from_xyz(math.pi * theta_xyz[i])\n",
    "    mesh_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(\n",
    "        size=1.0, origin=theta_xyz[i]\n",
    "    )\n",
    "    mesh_frame.transform(T_init)\n",
    "    vis.add_geometry(mesh_frame)\n",
    "vis.run()\n",
    "vis.destroy_window()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "def fit_line_ransac(points, num_iterations=100, threshold=0.1):\n",
    "    \"\"\"\n",
    "    Fits a line to a set of 2D points using RANSAC.\n",
    "    \n",
    "    Args:\n",
    "        points (list of tuples): List of (x, y) coordinates.\n",
    "        num_iterations (int): Number of RANSAC iterations.\n",
    "        threshold (float): Distance threshold for inliers.\n",
    "    \n",
    "    Returns:\n",
    "        tuple: (slope, intercept) of the fitted line.\n",
    "    \"\"\"\n",
    "    best_slope, best_intercept = None, None\n",
    "    best_inlier_count = 0\n",
    "    \n",
    "    for _ in range(num_iterations):\n",
    "        sample_indices = np.random.choice(len(points), 2, replace=False)\n",
    "        sample_points = [points[i] for i in sample_indices]\n",
    "        \n",
    "        x1, y1 = sample_points[0]\n",
    "        x2, y2 = sample_points[1]\n",
    "        \n",
    "        # Fit a line to the sample points\n",
    "        slope = (y2 - y1) / (x2 - x1)\n",
    "        intercept = y1 - slope * x1\n",
    "        \n",
    "        # Count inliers\n",
    "        inlier_count = sum(abs(y - (slope * x + intercept)) < threshold for x, y in points)\n",
    "        \n",
    "        if inlier_count > best_inlier_count:\n",
    "            best_inlier_count = inlier_count\n",
    "            best_slope, best_intercept = slope, intercept\n",
    "    \n",
    "    return best_slope, best_intercept\n",
    "\n",
    "# Example usage\n",
    "# Generate some noisy data points around a line\n",
    "true_slope, true_intercept = 2.0, 5.0\n",
    "noise_stddev = 0.5\n",
    "num_points = 100\n",
    "np.random.seed(42)\n",
    "points = [(x, true_slope * x + true_intercept + np.random.normal(0, noise_stddev)) for x in np.linspace(0, 10, num_points)]\n",
    "\n",
    "# Fit a line using RANSAC\n",
    "ransac_slope, ransac_intercept = fit_line_ransac(points)\n",
    "\n",
    "print(f\"True line: y = {true_slope:.2f}x + {true_intercept:.2f}\")\n",
    "print(f\"RANSAC line: y = {ransac_slope:.2f}x + {ransac_intercept:.2f}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
